from typing import Tuple
import torch
from torchvision.ops import box_iou
from torchvision.ops import batched_nms
from torch import Tensor

# def calculate_density(boxes, iou_threshold=0.5):
#     """
#     计算每个框与周围框的最大重叠度（密度）
#     """
#     B, N, _ = boxes.shape
#     density = torch.zeros(B, N, device=boxes.device)
#     density_indices = torch.zeros(B, N, device=boxes.device, dtype=torch.int64)

#     # 计算所有框的IoU矩阵
#     for b in range(B):
#         iou_matrix = box_iou(boxes[b], boxes[b])
        
#         # 将对角线上的IoU值置为0，排除自身
#         torch.diagonal(iou_matrix).zero_()
        
#         # 找到每个框的IoU值超过阈值的邻居的最大重叠度
#         valid_iou_mask = iou_matrix > iou_threshold
#         # 如果没有超过阈值的邻居，则使用0填充
#         density_scores = torch.where(valid_iou_mask, iou_matrix, torch.zeros_like(iou_matrix))
#         # 计算每行的最大值及对应的索引
#         density[b], density_indices[b] = torch.max(density_scores, dim=1)

#     return density, density_indices

def calculate_density(boxes, iou_threshold=0.5):
    """
    计算每个框与周围框的最大重叠度（密度）
    - boxes: 候选框坐标，格式为[B, N, 4]（x1, y1, x2, y2）
    - iou_threshold: IoU阈值
    """
    B, N, _ = boxes.shape
    
    # 计算所有框的IoU矩阵
    iou_matrix = torch.stack([box_iou(boxes[b], boxes[b]) for b in range(B)])
    
    # 创建一个掩码，用于将对角线置为0
    mask = ~torch.eye(N, dtype=torch.bool, device=boxes.device)
    iou_matrix = iou_matrix * mask.unsqueeze(0)
    
    # 找到每个框的IoU值超过阈值的邻居的最大重叠度
    valid_iou_mask = iou_matrix > iou_threshold
    # 如果没有超过阈值的邻居，则使用0填充
    density = torch.where(valid_iou_mask, iou_matrix, torch.zeros_like(iou_matrix))
    
    return density

    # # 计算每行的最大值及对应的索引
    # density, density_indices = torch.max(density_scores, dim=2)
    
    # return density, density_indices


def density_weighted_score(boxes, scores, iou_threshold=0.5, alpha=0.5):
    """
    基于密度加权的NMS算法：
    - boxes: 候选框坐标，格式为[B, N, 4]（x1, y1, x2, y2）
    - scores: 置信度，格式为[B, N]
    - iou_threshold: 传统NMS的IoU阈值
    - alpha: 置信度与定位质量的平衡系数（0~1）
    """
    # Step 1: 计算每个框的定位质量（此处以周围框密度为例）
    
    densities = calculate_density(boxes, iou_threshold)
    _, density_indices = torch.max(densities, dim=2)
    density_scores = torch.gather(scores, 1, density_indices)
    densities_sum = torch.sum(densities, dim=-1)
    densities_sum[densities_sum > 0] = 1
    density_scores = density_scores * densities_sum

    # Step 2: 计算综合评分（置信度 + 密度权重）
    # 定位质量越高（密度越大），权重越高
    combined_scores = alpha * scores + (1 - alpha) * density_scores
    
    return densities, combined_scores

    # # Step 3: 执行传统批量NMS，但基于新的排序
    # keep_indices = batched_nms(boxes, combined_scores, idxs, iou_threshold)

    # return keep_indices, combined_scores

def test_density_weighted_nms():    
    # 假设我们有两个批次，每个批次有两个检测框
    boxes = torch.tensor([
        [[10, 10, 50, 50],  # 框A，索引0，批次0
         [12, 12, 58, 58],  # 框A，索引0，批次0
         [25, 25, 68, 68]], # 框B，索引1，批次0
        [[20, 20, 70, 70],  # 框A，索引0，批次1
         [25, 25, 75, 75],  # 框A，索引0，批次0
         [30, 30, 80, 80]]  # 框B，索引1，批次1
    ], dtype=torch.float32)

    scores = torch.tensor([
        [0.8, 0.88, 0.9],  # 批次0
        [0.85, 0.9, 0.95]  # 批次1
    ])
    B = len(scores[-1])  # 批次数量
    idxs = torch.tensor([0, 0, 0, 1, 1, 1])  # 类别索引，每个框对应一个类别
    
    densities, weighted_scores = density_weighted_score(boxes, scores, iou_threshold=0.5, alpha=0.9)

    keep_indices = batched_nms(boxes.reshape(-1, 4), weighted_scores.reshape(-1), idxs, iou_threshold=0.2)

    print("保留的检测框索引:", keep_indices)
    print("original_scores:", scores.reshape(-1))
    print("weighted_scores:", weighted_scores.reshape(-1))
    
    keep_idxs = idxs[keep_indices]
    print("类别索引:", keep_idxs)
    
    for b, keep in zip(keep_idxs, keep_indices % B):
        print(f"批次 {b} 中保留的检测框索引:", keep, end=" ")
        print(f"中对应的检测框:", boxes[b, keep], end=" ")
        print(f"中对应的综合评分:", weighted_scores[b, keep])

def top_iou(boxes, peak_aff_batched):
    top4 = torch.topk(peak_aff_batched, k=4, dim=1)
    top4_values, top4_indices = top4.values, top4.indices
    top4_indices_expanded = top4_indices.unsqueeze(-1)
    top4_indices_expanded = top4_indices_expanded.expand(-1, -1, 4)
    top4_boxes = torch.gather(boxes, 1, top4_indices_expanded)
    top01iou = torch.diag(box_iou(top4_boxes[:, 0, :], top4_boxes[:, 1, :]))
    top12iou = torch.diag(box_iou(top4_boxes[:, 1, :], top4_boxes[:, 2, :]))
    topiou_r = torch.clamp( top01iou/ (top12iou + 1e-9), 0, 1)
    
    return topiou_r

if __name__ == '__main__':
    test_density_weighted_nms()
